#include <vector>
#include <map>
#include <set>
#include <cstring>
#include <stdexcept>
#include <iostream>
#include <algorithm>
#include "nn.h"

class fn {
    public:
    std::vector<data*> inputs; // free pointers to heap-allocated manually
    std::vector<data*> outputs;
    std::map<std::pair<int,int>,grad*> grads;
    virtual void set_inputs(std::vector<data*> inputs) = 0;
    virtual void set_grads(std::set<std::pair<int,int>>) = 0;
    virtual std::set<std::pair<int,int>> get_valid_grads() const = 0;
    void allocate_inputs() { for(int i = 0; i < inputs.size(); i ++) if(inputs[i]->vector == nullptr) inputs[i]->allocate(); }
    void allocate_outputs() { for(int i = 0; i < outputs.size(); i ++) if(outputs[i]->vector == nullptr) outputs[i]->allocate(); }
    void allocate_grads() { for(auto i = grads.begin(); i != grads.end(); i ++) if(i->second->vector == nullptr) i->second->allocate(); }
    ~fn() {
        for(int i = 0; i < outputs.size(); i ++) delete outputs[i];
        for(auto i = grads.begin(); i != grads.end(); i ++) delete i->second;
    }
    virtual void operator()() = 0;
};

class biased_lincomb : public fn {
    public:
    void operator()() override {
        data_linear* values = (data_linear*)(inputs[0]);
        data_matrix* coefficients = (data_matrix*)(inputs[1]);
        data_linear* biases = (data_linear*)(inputs[2]);
        data_linear* combinations = (data_linear*)(outputs[0]);
        coefficients->matmul(*values, *combinations);
        for(int i = 0; i < biases->length; i ++)
            combinations->vector[i] += biases->vector[i];
        if(grads.count({0, 0})) {
            //if(grads[{0, 0}]->type != grad_type::DENSE) throw std::runtime_error("dense_lincomb: Wrong type grad by values");
            dense_grad* values_grad = (dense_grad*)(grads[{0, 0}]);
            memcpy(values_grad->vector, coefficients->vector, coefficients->size() * sizeof(dtype));
        }
        if(grads.count({1, 0})) {
            //if(grads[{1, 0}]->type != grad_type::LINCOMB) throw std::runtime_error("dense_lincomb: Wrong type grad by coefficients");
            lincomb_grad* coefficients_grad = (lincomb_grad*)(grads[{1, 0}]);
            memcpy(coefficients_grad->vector, values->vector, values->size() * sizeof(dtype));
        }
        if(grads.count({2, 0})) {
            //if(grads[{2, 0}]->type != grad_type::SCALAR) throw std::runtime_error("dense_lincomb: Wrong type grad by biases");
            scalar_grad* biases_grad = (scalar_grad*)(grads[{2, 0}]);
            for(int i = 0; i < biases->length; i ++) biases_grad->vector[i] = 1;
        }
    }
    void set_inputs(std::vector<data*> inputs) override {
        if(inputs.size() != 3) throw std::runtime_error("dense_lincomb::set_inputs: Wrong number of inputs");
        if(inputs[0]->type != data_type::LINEAR) throw std::runtime_error("dense_lincomb::set_inputs: Wrong type of input 0");
        if(inputs[1]->type != data_type::MATRIX) throw std::runtime_error("dense_lincomb::set_inputs: Wrong type of input 1");
        if(inputs[2]->type != data_type::LINEAR) throw std::runtime_error("dense_lincomb::set_inputs: Wrong type of input 2");
        if(((data_linear*)inputs[0])->length != ((data_matrix*)inputs[1])->width) throw std::runtime_error("dense_lincomb::set_inputs: Input does not match with matrix");
        if(((data_matrix*)inputs[1])->height != ((data_linear*)inputs[2])->length) throw std::runtime_error("dense_lincomb::set_inputs: Wrong number of biases");
        this->inputs = inputs;
        if(!outputs.size() || ((data_linear*)outputs[0])->length != ((data_matrix*)inputs[1])->height) {
            if(outputs.size()) delete outputs[0];
            outputs = {new data_linear(((data_matrix*)inputs[1])->height)};
        }
    }
    void set_grads(std::set<std::pair<int,int>> links) override {
        for(auto i : grads)
            if(!links.count(i.first))
                delete i.second;
        for(auto i : links) {
            int inp = i.first, out = i.second;
            if(inp == 0 && out == 0)
                if(!grads.count(i) || grads[i]->type != grad_type::DENSE ||
                   ((dense_grad*)grads[i])->in_size != ((data_linear*)inputs[0])->length ||
                   ((dense_grad*)grads[i])->out_size != ((data_matrix*)inputs[1])->height) {
                    if(grads.count(i)) delete grads[i];
                    grads[i] = new dense_grad(((data_linear*)inputs[0])->length, ((data_matrix*)inputs[1])->height);
                }
            else if(inp == 1 && out == 0)
                if(!grads.count(i) || grads[i]->type != grad_type::LINCOMB ||
                   ((lincomb_grad*)grads[i])->vars != ((data_linear*)inputs[0])->length ||
                   ((lincomb_grad*)grads[i])->out_size != ((data_matrix*)inputs[1])->height) {
                    if(grads.count(i)) delete grads[i];
                    grads[i] = new dense_grad(((data_linear*)inputs[0])->length, ((data_matrix*)inputs[1])->height);
                }
            else if(inp == 2 && out == 0)
                if(!grads.count(i) || grads[i]->type != grad_type::SCALAR ||
                   ((scalar_grad*)grads[i])->in_size != ((data_matrix*)inputs[1])->height) {
                    if(grads.count(i)) delete grads[i];
                    grads[i] = new scalar_grad(((data_matrix*)inputs[1])->height);
                }
            else throw std::runtime_error("biased_lincomb::set_grads: Wrong link");
        }
    }
    // get_valid_grads = {{input, output}...}
    std::set<std::pair<int,int>> get_valid_grads() const override { return {{0, 0}, {1, 0}, {2, 0}}; }
};

struct graph_edge {
    int from_layer_p1, from_out; // taking out from_out of layer from_layer_p1-1 as input
    int to_layer_p1, to_out; // taking out to_out of layer to_layer_p1-1 as output
    int copy_number = 0; // 0 if single copy (or combination result); 1, 2, 3 ... if multiple copies
};

enum class algorythm_operation { COMBINE, SUM };
struct algorythm_node {
    graph_edge src1, src2;
    algorythm_operation operation;
    graph_edge dst;
};

class combined_fn : public fn {
    public:
    std::vector<fn*> layers;
    std::vector<std::vector<std::pair<int,int>>> connects; // connects[layer][input] = {layer + 1, output}
    std::vector<std::pair<int, int>> out_connects; // out_connects[total_output] = {layer + 1, output}
    std::map<graph_edge,grad*> internal_grads;
    std::vector<algorythm_node> algorythm;
    combined_fn(std::vector<fn*> layers, std::vector<std::vector<std::pair<int,int>>> connects, std::vector<std::pair<int,int>> out_connects) : layers(layers), connects(connects), out_connects(out_connects) {
        if(connects.size() != layers.size()) throw std::runtime_error("combined_fn: Invalid connects vector");
        for(int i = 0; i < layers.size(); i ++)
            for(std::pair<int,int> connect : connects[i])
                if(connect.first < 0 || connect.first > i) throw std::runtime_error("combined_fn: connect to wrong layer");
        for(std::pair<int,int> connect : out_connects)
            if(connect.first < 0 || connect.first > layers.size()) throw std::runtime_error("combined_fn: Taking output from wrong layer");
    }
    void set_inputs(std::vector<data*> inputs) override {
        std::vector<std::vector<data*>> data_layers;
        for(int i = 0; i < layers.size(); i ++) {
            std::vector<data*> layer_inputs;
            for(std::pair<int,int> connect : connects[i]) {
                std::vector<data*> data_layer = data_layers[connect.first];
                if(connect.second >= data_layer.size()) throw std::runtime_error("combined_fn::set_inputs: connect to wrong data");
                layer_inputs.push_back(data_layer[connect.second]);
            }
            layers[i]->set_inputs(layer_inputs);
            data_layers.push_back(layers[i]->outputs);
        }
        this->inputs = inputs;
        this->outputs.clear();
        for(int i = 0; i < out_connects.size(); i ++) {
            fn* layer = layers[out_connects[i].first];
            if(out_connects[i].second >= layer->outputs.size()) throw std::runtime_error("combined_fn::set_inputs: Taking not existing output");
            this->outputs.push_back(layer->outputs[out_connects[i].second]);
        }
    }
    // get_agorythm_backpropagation = {algorythm, graph}
    std::pair<std::vector<algorythm_node>,std::map<graph_edge,grad*>> get_agorythm_backpropagation(std::map<graph_edge,grad*> graph) {
        std::vector<algorythm_node> algorythm;
        for(int i = layers.size() - 1; i > 0; i --) { // reducing outs of layer i-1, ins of layer i
            std::map<int,std::pair<std::set<std::pair<graph_edge,grad*>>,std::set<std::pair<graph_edge,grad*>>>> nodes;
            for(std::pair<graph_edge,grad*> edge : graph) {
                if(edge.first.to_layer_p1 == layers.size() && edge.first.from_layer_p1 == i && !edge.first.copy_number)
                    nodes[edge.first.from_out].second.insert(edge);
                if(edge.first.to_layer_p1 == i && !edge.first.copy_number)
                    nodes[edge.first.to_out].first.insert(edge);
            }
            for(auto node : nodes) {
                std::set<std::pair<graph_edge,grad*>> edge1s = node.second.first;
                std::set<std::pair<graph_edge,grad*>> edge2s = node.second.second;
                std::set<std::pair<graph_edge,grad*>> edges;
                for(std::pair<graph_edge,grad*> edge1 : edge1s)
                    for(std::pair<graph_edge,grad*> edge2 : edge2s) {
                        graph_edge edge = graph_edge{edge1.first.from_layer_p1, edge1.first.from_out, edge2.first.to_layer_p1, edge2.first.to_out};
                        grad* g = edge1.second->get_combine_dst(*edge2.second);
                        graph.insert({edge, g});
                        algorythm.
                    }
            }
        }
    }
    void set_grads(std::set<std::pair<int,int>> crosslinks) override {
        auto linkss = gen_grad_graph(crosslinks);
        for(int i = 0; i < layers.size(); i ++) layers[i]->set_grads(linkss[i]);
        std::map<graph_edge,grad*> graph;
        for(int i = 0; i < layers.size(); i ++) {
            layers[i]->set_grads(linkss[i]);
            for(std::pair<int,int> link : linkss[i]) {
                std::pair<int,int> connect = connects[i][link.first];
                graph[graph_edge{connect.first, connect.second, i, link.second}] = layers[i]->grads[link];
            }
        }
        
    }
    protected:
    // grow_grad_graph_inp[layer] = {{input, out}...}
    std::vector<std::set<std::pair<int,int>>> grow_grad_graph_inp(std::set<int> ins) {
        std::vector<std::set<int>> outss = {ins}; // outss[layer + 1] = {used outputs}
        std::vector<std::set<std::pair<int,int>>> linkss;
        for(int k = 0; k < layers.size(); k ++) { // k = layer
            std::set<std::pair<int,int>> links; // {{input, out}...}
            std::set<int> datas; // {used output datas}
            for(int i = 0; i < connects[k].size(); i ++) { // i = input
                if(!outss[connects[k][i].first].count(connects[k][i].second)) continue;
                for(std::pair<int,int> link : layers[k]->get_valid_grads()) {
                    if(link.first == connects[k][i].second) {
                        datas.insert(link.second);
                        links.insert(link);
                    }
                }
            }
            outss.push_back(datas);
            linkss.push_back(links);
        }
        return linkss;
    }
    // grow_grad_graph_out[layer] = {{input, out}...}
    std::vector<std::set<std::pair<int,int>>> grow_grad_graph_out(std::set<std::pair<int,int>> outs) { // outs = {{layer + 1, out}...}
        std::vector<std::set<std::pair<int,int>>> linkss(layers.size(), std::set<std::pair<int,int>>()); // linkss[layer] = {{input, out}...}
        while(outs.size()) {
            std::pair<int,int> out = *outs.begin();
            outs.erase(out);
            for(std::pair<int,int> link : layers[out.first - 1]->get_valid_grads()) { // link = {input, output}
                if(link.second == out.second) {
                    std::pair<int,int> connect = connects[out.first - 1][out.second];
                    if(connect.first) {
                        outs.insert(connect);
                        linkss[out.first - 1].insert(link);
                    }
                }
            }
        }
        return linkss;
    }
    // gen_grad_graph[layer] = {{input, out}...}
    // crosslinks = {{total_input, total_output}...}
    std::vector<std::set<std::pair<int,int>>> gen_grad_graph(std::set<std::pair<int,int>> crosslinks) {
        std::vector<std::set<std::pair<int,int>>> linkss(layers.size(), std::set<std::pair<int,int>>());
        for(std::pair<int,int> crosslink : crosslinks) {
            auto linkss1 = grow_grad_graph_inp({crosslink.first});
            auto linkss2 = grow_grad_graph_out({out_connects[crosslink.second]});
            for(int i = 0; i < layers.size(); i ++)
                for(std::pair<int,int> link : linkss1[i])
                    if(linkss2[i].count(link))
                        linkss[i].insert(link);
        }
        return linkss;
    }
};

int main() {
    //std::cout << std::endl;
    //std::cout << typeid(lincomb_grad).name() << std::endl;
    lincomb_grad l(10, 10);
    l.get_combine_dst(dense_grad(10, 10));
    return 0;
}